#include "random_in.h"
#include "sm3.h"
#include "mpi.h"
#include <time.h>
#include <stdlib.h>

static internal_state_t in_state;

/**
 * @brief 比特串V从左边起的a个比特
 * 
 * @param V 
 * @param a 
 * @return int 
 */
int leftmost(unsigned char* V, const int a, unsigned char* out)
{
    int byte_len = a/8;
    int left_bits = a - 8*byte_len;

    memcpy(out, V, byte_len);

    if(left_bits > 1){
        unsigned char temp = 0xff >> (left_bits);
        temp = ~temp;
        out[byte_len] = temp & V[byte_len];
    }
    return 0;
}

int _get_entropy(unsigned char* entropy, const int len)
{
    unsigned char buf[SM3_DATA_LEN];
    unsigned char hash[SM3_DATA_LEN];
    sm3_ctx ctx;
    int ret = 0;

    ret = read_entropy_sys(buf, SM3_DATA_LEN);
    if(ret) return ret;

    sm3_init(&ctx);
    sm3_update(&ctx, buf, SM3_DATA_LEN);
#if USE_DEV_ENTROPY_SOURCE
    if(!read_entropy_dev(buf, min_entropy_bytes))
        sm3_update(&ctx, buf, min_entropy_bytes);
#endif
#if USE_CPU_ENTROPY_SOURCE
    if(!read_entropy_cpu(buf, min_entropy_bytes))
        sm3_update(&ctx, buf, min_entropy_bytes);
#endif
    sm3_final(&ctx, hash);
    memcpy(entropy, hash, len);
    return ret;
}

/**
 * @brief 从熵源获取一串比特
 *  要求返回的比特串至少含有 min_entropy 比特熵，
 *  长度大于等于 min_entropy_input_length，
 *  小于等于 max_entropy_input_length
 * 
 * @param entropy 
 * @param len 
 * @return int 
 */
int get_entropy(unsigned char* entropy, const long len)
{
    int ret = 0;
    int loop = len/SM3_DATA_LEN;
    int left = len - loop*SM3_DATA_LEN;

    if((len > max_entropy_input_length/8) ||
        (len < min_entropy_input_length/8))
        return ERANGE;

    for (int i = 0; i < loop; i++)
    {
        ret = _get_entropy(entropy+i*min_entropy_bytes, min_entropy_bytes);
        if(ret) return ret;
    }
    if(left > 0){
        ret = _get_entropy(entropy+loop*min_entropy_bytes, left);
    }
    return ret;
}

static inline int int_ceil(int a, int b)
{
    int t = a/b;
    if(a%b) t++;
    return t;
}

/**
 * @brief sm3派生函数，对输入字符串进行杂凑运算，
 *  返回长度为 number_of_bits_to_return 的比特串
 * 
 * @param str 
 * @param strlen 
 * @param out 
 * @return int 
 */
int sm3_df(const unsigned char* str, const int slen, unsigned char* out)
{
    // temp = NULL
    unsigned char temp[4096] ={0};
    uint32_t number_of_bits_to_return_ = number_of_bits_to_return;
    uint32_t counter = 0x01;
    sm3_ctx ctx;
    // len = [number_of_bytes_to_return/SM3_DATA_LEN]
    int len = int_ceil(number_of_bytes_to_return, SM3_DATA_LEN);

    for (int i = 0; i < len; i++)
    {
        // temp = temp||sm3(counter||number_of_bits_to_return||input_str)
        sm3_init(&ctx);
        sm3_update(&ctx, (const unsigned char*)&counter, sizeof(counter));
        sm3_update(&ctx, (const unsigned char*)&number_of_bits_to_return_, sizeof(uint32_t));
        sm3_update(&ctx, str, slen);
        sm3_final(&ctx, temp+i*SM3_DATA_LEN);

        // counter  = counter + 1
        counter++;
        if((i*SM3_DATA_LEN) > (4096-SM3_DATA_LEN))
            return ELNRNG;
    }
    // requested_bits = leftmost(temp, number_of_bits_to_return)
    return leftmost(temp, number_of_bits_to_return, out);
}

/**
 * @brief 初始化函数
 *  利用个性化字符串、noise, 对内部状态进行设置。
 * @param str 
 * @param slen 
 * @param noise 
 * @param nlen 
 * @return int 
 */
int sm3_rng_instantiate(const char* str, const int slen,
    const void* noise, const int nlen)
{
    int ret = 0;
    if(!noise || nlen < 1) return EINVAL;
    // entropy_input = get_entropy()
    unsigned char entropy_input[min_entropy_bytes];
    ret = get_entropy(entropy_input, min_entropy_bytes);
    if( ret ) return ret;

    // seed = sm3_df(entropy_input||noise||str)
    unsigned char* input = malloc(slen+nlen+min_entropy_bytes);
    if(!input) return errno;
    memcpy(input, entropy_input, min_entropy_bytes);
    memcpy(input+min_entropy_bytes, noise, nlen);
    if(str && slen > 1)
        memcpy(input+min_entropy_bytes+nlen, str, slen);

    unsigned char seed[SEEDLEN/8];
    ret = sm3_df(input, min_entropy_bytes+nlen+slen, seed);
    free(input);
    if(ret) return ret;

    // V = seed
    memcpy(in_state.V, seed, SEEDLEN_BYTES);

    // C = smd_df(0x00||V)
    unsigned char _c[SEEDLEN_BYTES+1];
    _c[0] = 0x00;
    memcpy(_c+1, in_state.V, SEEDLEN_BYTES);
    ret = sm3_df(_c, SEEDLEN_BYTES+1, in_state.C);
    if(ret) return ret;

    in_state.reseed_counter = 1;
    in_state.last_reseed_time = get_current_time_in_second();
    return ret;
}

/**
 * @brief 重播种函数
 * 在提供新种子值的情况下更新内部状态
 * @param entropy 熵源获得的比特串
 * @param elen 熵长度
 * @param opt 可选字符串
 * @param optlen 
 * @return int 
 */
int sm3_rng_reseed(const unsigned char* entropy, const int elen,
    const unsigned char* opt, const int optlen)
{
    unsigned char seed[SEEDLEN_BYTES] = {0};
    unsigned char prefix = 0x01;

    if(!entropy) return EINVAL;
    if(!opt && optlen != 0) return EINVAL;

    // sed = sm3_df(0x01||entropy_input||V||additional_input)
    int tatollen = 1+elen+SEEDLEN_BYTES+optlen;
    unsigned char* input = malloc(tatollen);
    if(!input) return errno;
    input[0] = prefix;
    memcpy(input+1, entropy, elen);
    memcpy(input+1+elen, in_state.V, SEEDLEN_BYTES);
    if(opt) memcpy(input+1+elen+SEEDLEN_BYTES, opt, optlen);
    sm3_df(input, tatollen, seed);
    // V = seed
    memcpy(in_state.V, seed, SEEDLEN_BYTES);
    // C = sm3_df(0x00||V)
    prefix = 0x00;
    input[0] = prefix;
    memcpy(input+1, in_state.V, SEEDLEN_BYTES);
    sm3_df(input, 1+SEEDLEN_BYTES, in_state.C);
    free(input);

    in_state.reseed_counter = 1;
    in_state.last_reseed_time = get_current_time_in_second();
    return 0;
}

/**
 * @brief 输出函数, 256bits(32byte)
 *  利用当前的内部状态和可能的额外输入返回随机数，并对内部状态进行更新。
 * @param opt 
 * @param optlen 
 * @param out 
 * @return int 
 */
int sm3_rng_generate( const void* opt, const int optlen,
    unsigned char out[32])
{
    int ret = 0;
    MPI _res= NULL, _v= NULL, _w= NULL, _sm= NULL;
    MPI __v = NULL, __h= NULL, __c= NULL, __rc= NULL, __sm= NULL;
    unsigned char* res_buf = NULL;
    sm3_ctx ctx;
    unsigned char prefix = 0x02;
    unsigned char hash[SM3_DATA_LEN] = {0};
    unsigned char m[SEEDLEN_BYTES+1] = {0};
    m[0] = 1;
    unsigned int nbyte;
    int sign;

    // 如果reseed_counter>reseed_interval_in_counter,
    // 或者Get_current_time_in_second()-last_reseed_time>reseed_interval_in_time，
    // 执行SM3_RNG_Reseed
    if(in_state.reseed_counter > reseed_interval_in_counter ||
        (get_current_time_in_second()-in_state.last_reseed_time >
        reseed_interval_in_time)) {
        unsigned char entropy[SEEDLEN_BYTES];
        ret = get_entropy(entropy, SEEDLEN_BYTES);
        if(ret) goto out;
        ret = sm3_rng_reseed(entropy, SEEDLEN_BYTES,
            (unsigned char*)opt, optlen);
        if(ret) goto out;
    }

    if(opt && optlen > 0){
        // W = sm3(0x02||V||additional_input)
        sm3_init(&ctx);
        sm3_update(&ctx, &prefix, sizeof(prefix));
        sm3_update(&ctx, in_state.V, SEEDLEN_BYTES);
        sm3_update(&ctx, opt, optlen);
        sm3_final(&ctx, hash);
        // V = (V+W) mod (2^seedlen)
        _res = mpi_new(0);
        _v = mpi_read_raw_data(in_state.V, SEEDLEN_BYTES);
        _w = mpi_read_raw_data(hash, SM3_DATA_LEN);
        _sm = mpi_read_raw_data(m, SEEDLEN_BYTES+1);
        if(!_res || !_v || !_w || !_sm){
            ret = ECANCELED;
            goto memout;
        }
        mpi_addm(_res, _v, _w, _sm);
        res_buf = mpi_get_buffer(_res, &nbyte, &sign);
        if(!res_buf){
            ret = ECANCELED;
            goto memout;
        }
        memcpy(in_state.V, res_buf, (nbyte>SEEDLEN_BYTES?SEEDLEN_BYTES:nbyte));
        free(res_buf);
        res_buf = NULL;
    }else{
        _res = mpi_new(0);
        _v = mpi_new(0);
        _w = mpi_new(0);
        _sm = mpi_new(0);
    }

    // temp = sm3(V)
    sm3(in_state.V, SEEDLEN_BYTES, hash);
    // returned_bits=leftmost(temp, requested_number_of_bits)；
    leftmost(hash, min_entropy_bits, out);
    // H = sm3(0x03||V)
    sm3_init(&ctx);
    prefix = 0x03;
    sm3_update(&ctx, &prefix, sizeof(prefix));
    sm3_update(&ctx, in_state.V, SEEDLEN_BYTES);
    sm3_final(&ctx, hash);
    // V = (V+H+C+reseed_counter)mod 2^seedlen
    __v = mpi_read_raw_data(in_state.V, SEEDLEN_BYTES);
    __h = mpi_read_raw_data(hash, SM3_DATA_LEN);
    __c = mpi_read_raw_data(in_state.C, SEEDLEN_BYTES);
    __rc = mpi_read_raw_data(&in_state.reseed_counter, sizeof(in_state.reseed_counter));
    __sm = mpi_read_raw_data(m, SEEDLEN_BYTES+1);

    if(!__v || !__h || !__c || !__rc || !__sm){
        ret = ECANCELED;
        goto memout2;
    }

    mpi_add(_v, __v, __h);
    mpi_add(_w, _v, __c);
    mpi_addm(_res, _w, __rc, __sm);
    res_buf = mpi_get_buffer(_res, &nbyte, &sign);
    if(!res_buf){
        ret = ENOMEM;
        goto memout2;
    }
    memcpy(in_state.V, res_buf, (nbyte>SEEDLEN_BYTES?SEEDLEN_BYTES:nbyte));
    free(res_buf);
    // reseed_counter=reseed_counter+1
    in_state.reseed_counter++;

memout2:
    mpi_free(__v);
    mpi_free(__h);
    mpi_free(__c);
    mpi_free(__rc);
    mpi_free(__sm);

memout:
    mpi_free(_res);
    mpi_free(_v);
    mpi_free(_w);
    mpi_free(_sm);
out:
    return ret;
}

long get_current_time_in_second()
{
    return time(NULL);
}